import os
import matplotlib.pyplot as plt
from typing import Dict, List, Tuple, Any, Optional
from rainforeLearn.gomoku.v2.train.constants.train_stats_printer_constants import TrainingStatsPrinterConstants

# 设置中文字体（适用于macOS）
plt.rcParams['font.sans-serif'] = ['PingFang SC', 'Arial Unicode MS', 'SimHei']
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


class ChartGenerator:
    """图表生成器"""
    
    def __init__(self, stats_data: Dict[str, Any]):
        self.stats_data = stats_data
    
    def generate_comprehensive_analysis_chart(self, episode: int, reward_summary: Optional[Dict] = None) -> str:
        """生成综合分析图表"""
        if not reward_summary:
            print("⚠️ 无奖励系统统计数据，跳过图表绘制")
            return ""
        
        batch_win_rates = reward_summary.get('recent_batch_win_rates', [])
        if len(batch_win_rates) < 3:
            print("⚠️ 批次胜率数据不足，跳过图表绘制")
            return ""
        
        fig, axes = self._create_figure_and_axes(episode)
        
        # 添加批次胜率变化趋势图
        self._plot_batch_win_rate_trend(axes[0, 0], batch_win_rates)
        # 添加总体胜率变化趋势图
        self._plot_overall_win_rate_trend(axes[0, 1])
        # 添加损函数失监控图表
        self._plot_loss_monitoring(axes[0, 2])
        # 添加惩罚系数图表
        self._plot_penalty_coefficient_trend(axes[1, 0], reward_summary)  
        # 添加威胁分布图表
        self._plot_threat_type_distribution(axes[1, 1])
        # 添加奖励惩罚累计图图
        self._plot_reward_penalty_cumulative(axes[1, 2])
        
        return self._save_chart(episode)
    
    def _create_figure_and_axes(self, episode: int) -> Tuple[plt.Figure, Any]:
        """创建图表和坐标轴"""
        fig, axes = plt.subplots(2, 3, figsize=TrainingStatsPrinterConstants.CHART_FIGURE_SIZE)
        fig.suptitle(f'五子棋训练综合分析 (Episode {episode})', fontsize=16, fontweight='bold')
        return fig, axes
    
    def _plot_batch_win_rate_trend(self, ax, batch_win_rates: List[float]) -> None:
        """绘制批次胜率趋势图"""
        batch_nums = range(1, len(batch_win_rates) + 1)
        ax.plot(batch_nums, batch_win_rates, 'b-', 
               linewidth=TrainingStatsPrinterConstants.CHART_LINE_WIDTH, marker='o')
        ax.set_title('批次胜率趋势')
        ax.set_xlabel('批次')
        ax.set_ylabel('胜率')
        ax.grid(True, alpha=0.3)
        ax.axhline(y=0.5, color='r', linestyle='--', alpha=0.5, label='50%基准线')
        ax.legend()
    
    def _plot_overall_win_rate_trend(self, ax) -> None:
        """绘制总体胜率变化趋势"""
        win_rate_history = self.stats_data.get('win_rate_history', [])
        
        if len(win_rate_history) > 1:
            episodes = [x[0] for x in win_rate_history]
            win_rates = [x[2] for x in win_rate_history]
            ax.plot(episodes, win_rates, 'g-', 
                   linewidth=TrainingStatsPrinterConstants.CHART_LINE_WIDTH, marker='o')
            ax.set_title('总体胜率变化趋势')
            ax.set_xlabel('Episode')
            ax.set_ylabel('胜率')
            ax.axhline(y=0.5, color='r', linestyle='--', alpha=0.5)
            ax.grid(True, alpha=0.3)
        else:
            self._plot_no_data_message(ax, '总体胜率变化趋势')

    def _plot_loss_monitoring(self, ax) -> None:
        """绘制损失函数值监控图表"""
        # 从stats_data中获取损失历史数据
        loss_history = self.stats_data.get('loss_history', [])

        if len(loss_history) > 0:
            # 计算最近100个数据点的统计信息
            loss_list = list(loss_history)
            recent_data = loss_list[-100:] if len(loss_list) >= 100 else loss_list

            # 创建x轴数据（最近的训练步数）
            x_data = list(range(len(recent_data)))

            # 绘制损失曲线
            ax.plot(x_data, recent_data, 'b-', alpha=0.7,
                    linewidth=1, label='Loss')

            # 计算并显示统计指标
            avg_loss = sum(recent_data) / len(recent_data)
            max_loss = max(recent_data)
            min_loss = min(recent_data)

            # 添加统计线
            ax.axhline(y=avg_loss, color='red', linestyle='--',
                       alpha=0.8, label=f'平均损失: {avg_loss:.4f}')
            ax.axhline(y=max_loss, color='orange', linestyle=':',
                       alpha=0.6, label=f'最大损失: {max_loss:.4f}')
            ax.axhline(y=min_loss, color='green', linestyle=':',
                       alpha=0.6, label=f'最小损失: {min_loss:.4f}')

            ax.set_title('损失监控 (最近100步)')
            ax.set_xlabel('训练步数')
            ax.set_ylabel('损失值')
            ax.legend(loc='upper right', fontsize=8)
            ax.grid(True, alpha=0.3)

            # 设置y轴范围，避免异常值影响显示
            y_range = max_loss - min_loss
            if y_range > 0:
                ax.set_ylim(min_loss - y_range * 0.1, max_loss + y_range * 0.1)
        else:
            self._plot_no_data_message(ax, '损失监控')

    def _plot_penalty_coefficient_trend(self, ax, reward_summary: Dict) -> None:
        """绘制惩罚系数变化趋势图"""
        coefficient_history = reward_summary.get('coefficient_history', [])
        
        if len(coefficient_history) > 1:
            # coefficient_history 格式: [(batch_num, coefficient_value), ...]
            batch_nums = [x[0] for x in coefficient_history]
            coefficients = [x[1] for x in coefficient_history]
            
            ax.plot(batch_nums, coefficients, 'purple', 
                   linewidth=TrainingStatsPrinterConstants.CHART_LINE_WIDTH, 
                   marker='o', markersize=4)
            ax.set_title('惩罚系数动态调整趋势')
            ax.set_xlabel('批次')
            ax.set_ylabel('惩罚系数')
            ax.grid(True, alpha=0.3)
            
            # 添加当前系数值标注
            if coefficients:
                current_coeff = coefficients[-1]
                ax.axhline(y=current_coeff, color='red', linestyle='--', 
                          alpha=0.5, label=f'当前系数: {current_coeff:.3f}')
                ax.legend()
        else:
            self._plot_no_data_message(ax, '惩罚系数动态调整趋势')
    
    def _plot_threat_type_distribution(self, ax) -> None:
        """绘制威胁类型分布"""
        threat_stats = self.stats_data.get('threat_defense_stats', {})
        threat_types = ['关键', '高', '中', '低']
        threat_counts = [
            threat_stats.get('critical_threats', 0),
            threat_stats.get('high_threats', 0),
            threat_stats.get('medium_threats', 0),
            threat_stats.get('low_threats', 0)
        ]
        
        if sum(threat_counts) > 0:
            ax.pie(threat_counts, labels=threat_types, autopct='%1.1f%%',
                  colors=['red', 'orange', 'yellow', 'lightblue'])
            ax.set_title('威胁类型分布')
        else:
            self._plot_no_data_message(ax, '威胁类型分布')
    
    def _plot_reward_penalty_cumulative(self, ax) -> None:
        """绘制奖励惩罚累计图"""
        win_rate_bonus_history = self.stats_data.get('win_rate_bonus_history', [])
        batch_penalty_history = self.stats_data.get('batch_penalty_history', [])
        
        events_by_episode = self._aggregate_events_by_episode(
            win_rate_bonus_history, batch_penalty_history)
        
        if events_by_episode:
            plot_data = self._calculate_cumulative_values(events_by_episode)
            self._plot_cumulative_lines(ax, plot_data)
        else:
            self._plot_no_data_message(ax, '奖励惩罚累计变化')
    
    def _aggregate_events_by_episode(self, bonus_history: List, penalty_history: List) -> Dict:
        """按episode聚合事件"""
        events_by_episode = {}
        
        for ep, reward_val in bonus_history:
            if ep not in events_by_episode:
                events_by_episode[ep] = {'reward': 0, 'penalty': 0}
            events_by_episode[ep]['reward'] += reward_val
        
        for ep, penalty_val in penalty_history:
            if ep not in events_by_episode:
                events_by_episode[ep] = {'reward': 0, 'penalty': 0}
            events_by_episode[ep]['penalty'] += penalty_val
        
        return events_by_episode
    
    def _calculate_cumulative_values(self, events_by_episode: Dict) -> Dict:
        """计算累计值"""
        sorted_episodes = sorted(events_by_episode.keys())
        plot_data = {
            'episodes': [],
            'cumulative_rewards': [],
            'cumulative_penalties': [],
            'cumulative_net': []
        }
        
        current_cum_reward = 0
        current_cum_penalty = 0
        
        for ep in sorted_episodes:
            current_cum_reward += events_by_episode[ep].get('reward', 0)
            current_cum_penalty += events_by_episode[ep].get('penalty', 0)
            
            plot_data['episodes'].append(ep)
            plot_data['cumulative_rewards'].append(current_cum_reward)
            plot_data['cumulative_penalties'].append(current_cum_penalty)
            plot_data['cumulative_net'].append(current_cum_reward - current_cum_penalty)
        
        return plot_data
    
    def _plot_cumulative_lines(self, ax, plot_data: Dict) -> None:
        """绘制累计线图"""
        ax.plot(plot_data['episodes'], plot_data['cumulative_rewards'], 
               'g-', label='累计奖励', linewidth=TrainingStatsPrinterConstants.CHART_LINE_WIDTH)
        ax.plot(plot_data['episodes'], plot_data['cumulative_penalties'], 
               'r-', label='累计惩罚', linewidth=TrainingStatsPrinterConstants.CHART_LINE_WIDTH)
        ax.plot(plot_data['episodes'], plot_data['cumulative_net'], 
               'b-', label='净收益', linewidth=TrainingStatsPrinterConstants.CHART_LINE_WIDTH)
        ax.set_title('奖励惩罚累计变化')
        ax.set_xlabel('Episode')
        ax.set_ylabel('累计分数')
        ax.legend()
        ax.grid(True, alpha=0.3)
    
    def _plot_no_data_message(self, ax, title: str) -> None:
        """绘制无数据消息"""
        ax.text(0.5, 0.5, '数据不足', ha='center', va='center', transform=ax.transAxes)
        ax.set_title(title)
    
    def _save_chart(self,episode: int) -> str:
        """保存图表"""
        os.makedirs(TrainingStatsPrinterConstants.ANALYSIS_CHARTS_DIR, exist_ok=True)
        chart_path = (f'{TrainingStatsPrinterConstants.ANALYSIS_CHARTS_DIR}/'
                     f'{TrainingStatsPrinterConstants.CHART_FILE_PREFIX}{episode}.png')
        
        plt.tight_layout()
        plt.savefig(chart_path, dpi=TrainingStatsPrinterConstants.CHART_DPI, bbox_inches='tight')
        print(f"📊 综合分析图表已保存: {chart_path}")
        plt.close()
        return chart_path